import uuid
from typing import List, Dict

from pymilvus import Collection, FieldSchema, DataType, CollectionSchema

from bot.openai_bot import OpenAIBot
from conf.config import logger
from db.milvus.milvus_db import MilvusClient


class DecorationTitleDAO(object):
    """ 家装标题访问对象 """

    def __init__(self, *args, **kwargs):
        self.alias = str(uuid.uuid4())  # 给本次的milvus链接起个名字，后续释放链接要用到该名字
        self.collection_name = "decoration_title"
        super(DecorationTitleDAO, self).__init__(*args, **kwargs)

    def create_milvus_conn(self):
        """ 创建milvus连接 """
        milvus_client = MilvusClient(alias=self.alias)
        return milvus_client

    def get_ids(self, title_vector: List[float], top_k: int) -> List[Dict]:
        """
        获取数据项的id列表、向量距离列表
        :param title_vector: 要查找的标题向量
        :param top_k: 返回前几个
        :return:
        """
        milvus_client = self.create_milvus_conn()
        collection = Collection(self.collection_name, using=self.alias)  # collection.load()  # 提前加载一次即可

        try:
            results = collection.search(
                data=[title_vector, ],
                anns_field="title_vector",
                param={"metric_type": "IP"},
                limit=top_k
            )
            print(results)
            ids = results[0].ids
            distances = results[0].distances

            ls = list()
            for i in range(len(ids)):
                ls.append({
                    "title_id": ids[i],
                    "distance": distances[i],
                })
            return ls
        finally:
            milvus_client.disconnect()

    def create_collection(self, dim: int):
        """
        创建集合
        :param dim: 矢量维度
        :return:
        """
        milvus_client = self.create_milvus_conn()
        try:

            fields = [
                FieldSchema(name="title_id", dtype=DataType.INT64, is_primary=True, auto_id=False, description="标题id"),
                FieldSchema(name="title_vector", dtype=DataType.FLOAT_VECTOR, dim=dim, description="标题对应的向量"),
            ]
            schema = CollectionSchema(fields=fields, description='家装标题')
            collection = Collection(name=self.collection_name, schema=schema, using=self.alias)

            index_params = {
                'metric_type': "IP",
                'index_type': "FLAT",
                'params': {}
            }
            collection.create_index(field_name='title_vector', index_params=index_params)
        finally:
            if "milvus_client" in dir():
                milvus_client.disconnect()


def main():
    openai_bot = OpenAIBot()
    vector = openai_bot.embeddings(text_list=["墙鼓包是怎么造成的", ])[0]
    baike_title = DecorationTitleDAO()
    res = baike_title.get_ids(title_vector=vector, top_k=5)
    logger.info(f"{res}")


if __name__ == '__main__':
    main()
